#include "parser.h"

int RegDistributer::GetRegisterBySym(string sym)
{
    int curr_reg_num = symbols_.size();
    for (int i = 0; i < curr_reg_num; i++) {
        if (sym == symbols_[i]) {
            return i;
        }
    }
    symbols_.emplace_back(sym);
    curr_reg_pos_++;
    curr_reg_num++;
    max_reg_num_ = max_reg_num_ > curr_reg_num ? max_reg_num_ : curr_reg_num;
    return curr_reg_num - 1;
}

int RegDistributer::AllocTempReg()
{
    curr_reg_pos_++;
    max_reg_num_ = max_reg_num_ > curr_reg_pos_ ? max_reg_num_ : curr_reg_pos_;
    return curr_reg_pos_ - 1;
}

InstSet::InstSet(string file_name)
{
    reg_dis_ = new RegDistributer();
    output_file_name_ = file_name;
}
InstSet::~InstSet()
{
    delete reg_dis_;
    for (auto byte_code : sets_) {
        delete byte_code;
    }
}

void InstSet::AddInst(ByteCode* inst)
{
    sets_.emplace_back(inst);
}

bool InstSet::WriteInstToFile()
{
    ofstream fout;
    fout.open(output_file_name_);
    if (!fout.is_open()) {
        cout << TAG << "Failed to open file " << output_file_name_ << "for write instructions" << endl;
    }
    fout << FOREST_VERSION << endl; // write forest compiler version to file for interpreter check
    fout << reg_dis_->GetRegisterNum() << endl; // write used register number
    for (auto code : sets_) {
        fout << code->GetCode(); // write byte code
    }
}

void InstSet::Print()
{
    for (auto code : sets_) {
        cout << code->GetCode(); // write byte code
    }
}

RegDistributer& InstSet::GetRegDis()
{
    return *reg_dis_;
}

Node::~Node()
{
    for (auto node : nodes_) {
        delete node;
    }
}

void SymbolNode::GenerateCode(InstSet& inst_set)
{
    int target = inst_set.GetRegDis().GetRegisterBySym(symbol_name_);
    int source_reg = inst_set.GetRegDis().AllocTempReg();
    inst_set.AddInst(new ByteCode(BYTE_CODE::MOV, target, source_reg));
}

void ValNode::GenerateCode(InstSet& inst_set)
{
    RegDistributer& reg_dis = inst_set.GetRegDis();
    int temp_reg = reg_dis.AllocTempReg();
    inst_set.AddInst(new ByteCode(BYTE_CODE::MOV_CON, temp_reg, val_));
    reg_dis.ReleaseTempReg();
}

void ExprNode::GenerateCode(InstSet& inst_set)
{
    RegDistributer& reg_dis = inst_set.GetRegDis();
    reg_dis.GetRegisterBySym(result_node_->ToString());
    cal_node_->GenerateCode(inst_set);
    result_node_->GenerateCode(inst_set);
}

void OpNode::GenerateCode(InstSet& inst_set)
{
    RegDistributer& reg_dis = inst_set.GetRegDis();
    int left_reg = reg_dis.AllocTempReg();
    int right_reg = reg_dis.AllocTempReg();
    inst_set.AddInst(new ByteCode(BYTE_CODE::MOV, left_reg, reg_dis.GetRegisterBySym(left_->ToString())));
    inst_set.AddInst(new ByteCode(BYTE_CODE::MOV, right_reg, reg_dis.GetRegisterBySym(right_->ToString())));
    inst_set.AddInst(new ByteCode(GetOpcodeByToken(token_), left_reg, right_reg));
    reg_dis.ReleaseTempReg();
    reg_dis.ReleaseTempReg();
}

BYTE_CODE OpNode::GetOpcodeByToken(TOKEN token)
{
    switch (token)
    {
    case TOKEN::ADD_TOKEN:
        return BYTE_CODE::ADD;
    case TOKEN::MINUS_TOKEN:
        return BYTE_CODE::MINUS;
    case TOKEN::MUL_TOKEN:
        return BYTE_CODE::MUL;
    case TOKEN::DIV_TOKEN:
        return BYTE_CODE::DIV;
    default:
        cout << TAG << __func__ << " unknown op token" << endl;
        return BYTE_CODE::UNKNOWN;
    }
}

void PrintNode::GenerateCode(InstSet& inst_set)
{
    RegDistributer& reg_dis = inst_set.GetRegDis();
    int reg = reg_dis.GetRegisterBySym(var_->ToString());
    inst_set.AddInst(new ByteCode(BYTE_CODE::PRINT, reg));
}

Parser::Parser(string input, string output)
{
    lexer_ = new Lexer(input);
    inst_set_ = new InstSet(output);
    cout << TAG << "Lexer constructed input " << input << endl;
}

Parser::~Parser()
{
    if (curr_token_ != nullptr) {
        delete curr_token_;
    }
    if (abs_syntac_tree_ != nullptr) {
        delete abs_syntac_tree_;
    }
    delete lexer_;
    delete inst_set_;
}

void Parser::Match(TOKEN token)
{
    if (DEBUG) {
        cout << TAG << "Current token type " << static_cast<int>(curr_token_->GetTokenType()) <<
            ", current token " << curr_token_->ToString() << endl;
    }
    if (curr_token_->GetTokenType() != token) {
        cout << TAG << "Failed to match token " << static_cast<int>(token) << ", current token is " << curr_token_->ToString() << endl;
        exit(-1);
    }
    curr_token_ = lexer_->GetNextToken();
}

void Parser::Expr()
{
    Node* symbol_node = new SymbolNode(curr_token_->ToString());
    Match(TOKEN::VAR_TOKEN);
    Match(TOKEN::ASSIGN_TOKEN);
    Node* val_node;
    if (curr_token_->GetTokenType() == TOKEN::VAR_TOKEN) {
        Node* op_left_node = new SymbolNode(curr_token_->ToString());
        Match(TOKEN::VAR_TOKEN);
        TOKEN token = curr_token_->GetTokenType();
        Match(token);
        Node* op_right_node = new SymbolNode(curr_token_->ToString());
        Match(TOKEN::VAR_TOKEN);
        Match(TOKEN::SEMI_TOKEN);
        val_node = new OpNode(token, op_left_node, op_right_node);
    } else {
        val_node = new ValNode(atoi(curr_token_->ToString().c_str()));
        Match(TOKEN::INT_TOKE);
        Match(TOKEN::SEMI_TOKEN);
    }
    Node* expr = new ExprNode(symbol_node, val_node);
    expr->GenerateCode(*inst_set_);
}

void Parser::Print()
{
    Match(TOKEN::PRINT_TOKEN);
    Node* symbolNode = new SymbolNode(curr_token_->ToString());
    Match(TOKEN::VAR_TOKEN);
    Match(TOKEN::SEMI_TOKEN);
    Node* print = new PrintNode(symbolNode);
    print->GenerateCode(*inst_set_);
}

void Parser::Parse()
{
    curr_token_ = lexer_->GetNextToken();
    while(curr_token_->GetTokenType() != TOKEN::END_TOKEN) {
        if (curr_token_->GetTokenType() == TOKEN::PRINT_TOKEN) {
            Print();
        } else {
            Expr();
        }
    }
    inst_set_->WriteInstToFile();
}